import subprocess
import itertools
import threading
import time
import logging

from tqdm import tqdm

from config import PROJECT_ROOT, hap_type, experiment_name, scenarios, rounds, target_temperature, trace_tag_mode
from analyze import analyze

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

hap_names = ["com.example.myapplication250227", "com.acts.ace.napitest", "com.rnoh.tester"]
center_x_position = 600
navigate_y_positions = {"green": 1400, "blue": 350}
scenario_y_positions = {
    "green": {
        "view": 420,
        "text": 1280,
        "image": 2100,
    },
    "blue": {"view": 220, "text": 440, "image": 660},
}


class CaptureTask(threading.Thread):

    def __init__(
        self,
        prefix: str = "",
        experiment_name: str = "",
        trace_tag_mode: str = "usual",
        is_rooted: bool = False,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.process = None

        self.prefix = prefix
        self.experiment_name = experiment_name
        self.trace_tag_mode = trace_tag_mode
        self.is_rooted = is_rooted

        self.trace_started = False
        self.trace_ended = False

    def run(self):
        self.process = subprocess.Popen(
            [
                "bash",
                "cap_hi.sh",
                self.prefix,
                self.experiment_name,
                "4" if self.is_rooted else "3",
                self.trace_tag_mode,
            ],
            cwd=PROJECT_ROOT / "script",
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )

        # Watch output for "tracing" keyword
        while True:
            output = self.process.stdout.readline()
            print(output)
            if output == "" and self.process.poll() is not None:
                self.trace_ended = True
                break
            if "tracing" in output.lower():
                self.trace_started = True
                break


def exec(cmd: list[str]):
    # print(" ".join(cmd))
    process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        bufsize=1,
        universal_newlines=True,
    )
    all_output = ""
    while True:
        output = process.stdout.readline()
        if output == "" and process.poll() is not None:
            break
        if output:
            all_output += output
            # print(output.strip())
    return_code = process.poll()
    if return_code != 0:
        raise subprocess.CalledProcessError(return_code, cmd)

    return all_output


def hdc(cmd: str, device: str = None):
    if isinstance(cmd, str):
        cmd = cmd.strip().split()
    if device:
        cmd = ["hdc", "-t", device, "shell"] + cmd
    else:
        cmd = ["hdc", "shell"] + cmd
    return exec(cmd)


def wait_for_device():
    logging.info("waiting for device")
    exec(["hdc", "wait-for-device"])
    for _ in range(300):
        if get_temp() is not None:
            break
        logging.error("execute hdc command failed. is debug permission granted?")
        time.sleep(1)


def get_tested_bundle_name():
    haps = hdc("bm dump -a")
    found = []
    for hap_name in hap_names:
        if hap_name in haps:
            found.append(hap_name)
    assert len(found) == 1, f"found {len(found)} testee haps"
    return found[0]


def go_home():
    hdc("uitest uiInput fling 600 2700 800 2000 20000")


def clear_bg_app(
    times: int = 1,
    ensure_cleared_haps: list[str] = [
        "com.rnoh.tester",
        "com.acts.ace.napitest",
        "com.example.myapplication250227",
        "com.huawei.hmos.calculator",
    ],
):
    go_home()
    for _ in range(10):
        hdc("uitest uiInput fling 600 2700 600 1600 2000")  # 上滑展示后台应用
        time.sleep(0.3)
        for _ in range(times):
            hdc("uitest uiInput fling 600 2000 600 1500 40000")  # 上滑清理
        time.sleep(0.1)

        need_clear = False
        for hap in ensure_cleared_haps:
            if hdc(f"pgrep {hap[-15:]}").strip() != "":
                need_clear = True
                break
        if not need_clear:
            break


def kill_app(bundle_name: str):
    hdc(f"aa stop-service -a EntryAbility -b {bundle_name}")


def get_temp():
    lines = hdc("hidumper -s BatteryService -a -i ").split("\n")
    for l in lines:
        if "temperature" in l:
            return int(l.split(":")[1].strip()) / 10


def get_app_vendor(bundle_name: str):
    info = hdc(f"bm dump -n {bundle_name} | grep vendor")
    if "vendor" not in info:
        return None
    info = info.splitlines()[-1]
    right = info.rfind('"')
    left = info.rfind('"', 0, right - 1) + 1
    return info[left:right]


def get_rom_version():
    info = hdc("param get const.product.software.version")
    # ALN-AL00 5.0.0.22(SP35DEVC00E22R4P1log)
    return info.strip().replace(".", "").replace("(", "_").replace(")", "").replace(" ", "_")


def get_udid(length: int = 8) -> str:
    return hdc("bm get --udid").split("\n")[1].strip()[:length]


def get_rooted() -> bool:
    return hdc("whoami").strip() == "root"


def get_phone_model() -> str:
    info = hdc("param get const.product.name")
    return info.strip().replace(" ", "_")


def start_app(bundle_name: str, ability_name: str = "EntryAbility"):
    hdc(f"aa start -a {ability_name} -b {bundle_name}")


def unlock():
    start_app("com.huawei.hmos.calculator", "CalculatorAbility")


def ensure_temperature(target_temperature: float, max_wait: int = 600):
    tic = time.time()
    temp = 0
    for idx in range(1300):
        temp = get_temp()
        if temp < target_temperature - 0.1:
            logging.info(f"temperature current: {temp}, target: {target_temperature}")
            unlock()
            time.sleep(1)
            clear_bg_app(1)
        else:
            break

    hot_last_time = False

    for idx in range(max_wait):
        if idx != 0:
            temp = get_temp()
        if temp <= target_temperature + 0.1:
            if hot_last_time:
                still_hot = False
                for _ in range(5):
                    time.sleep(1)
                    temp = get_temp()
                    if temp > target_temperature + 0.1:
                        still_hot = True
                        break
                hot_last_time = still_hot
                continue
            break
        logging.info(f"temperature current: {temp}, target: {target_temperature}")
        time.sleep(5)
        hot_last_time = True
    held_up = time.time() - tic
    if held_up > 1:
        logging.info(f"held up by temperature for {held_up}")
    if hot_last_time:
        raise RuntimeError("Temperature is too high, stop testing")
    held_up = held_up > 1
    
    # screen may be locked by now
    if held_up:
        unlock()
        time.sleep(1)
        clear_bg_app(1)
    
    return held_up


def run_benchmark(
    bundle_name: str,
    hap_type: str,
    experiment_name: str,
    scenario: str,
    target_temperature: float,
    trace_tag_mode: str,
    should_record: bool = True,
    post_sleep: float = 0,
):
    is_rooted = get_rooted()
    if not is_rooted and trace_tag_mode == "load":
        raise RuntimeError("Cannot capture load info on non-rooted device")

    ensure_temperature(target_temperature)

    if should_record:
        t = CaptureTask(
            prefix=f"{scenario}@{get_app_vendor(bundle_name)}@{experiment_name}@{get_rom_version()}@{"rooted" if is_rooted else "notrooted"}@{target_temperature}@{get_phone_model()}@{get_udid()}",
            experiment_name=experiment_name,
            trace_tag_mode=trace_tag_mode,
            is_rooted=is_rooted,
        )
        t.start()

    start_app(bundle_name)

    time.sleep(0.5)

    hdc(f"uitest uiInput click {center_x_position} {navigate_y_positions[hap_type]}")
    # green hap navigates to another page here, wait a bit
    if hap_type == "green":
        time.sleep(0.3)

    if should_record:
        started = False
        for _ in range(50):
            if t.trace_started:
                started = True
                break
            time.sleep(0.1)
        if not started:
            return
    # even when tracing is outputted recording doesn't actually start, wait a bit
    time.sleep(1)
    # trace recording takes longer to start when rooted
    if is_rooted:
        time.sleep(1.5)

    # print("======= clicking =========")
    hdc(f"uitest uiInput click {center_x_position} {scenario_y_positions[hap_type][scenario]}")

    # hands free during rendering
    time.sleep(2.5)
    # can sleep longer in some cases, like for 5.0 image
    time.sleep(post_sleep)
    clear_bg_app(1)

    # wait for backgroup process to terminate, else next record would fail
    if should_record:
        t.join(10)


def run_test(bundle_name, hap_type, experiment_name, scenarios, rounds, target_temperature, trace_tag_mode):
    logging.info("device connected")

    # unlock device
    unlock()
    clear_bg_app(10)

    ensure_temperature(target_temperature)

    for scenario, idx in tqdm(list(itertools.product(scenarios, range(rounds)))):
        print(scenario, idx)
        run_benchmark(
            bundle_name,
            hap_type,
            experiment_name,
            scenario,
            target_temperature,
            trace_tag_mode,
            True,
            # 20 if scenario == "image" and idx == 0 else 0,
            0,
        )
        time.sleep(1)

    clear_bg_app(1)
    go_home()


# """test clear background app"""
# fail = 0
# bundle_name = get_tested_bundle_name()
# for _ in tqdm(range(100)):
#     start_app(bundle_name)
#     time.sleep(0.5)
#     clear_bg_app(1)
#     time.sleep(1)
#     if hdc(f"pgrep {bundle_name}").strip() != "":
#         fail += 1
#         print("failed")
# print(f"fail: {fail}")
# exit()


if __name__ == "__main__":
    wait_for_device()

    # tested hap bundle name
    bundle_name = get_tested_bundle_name()

    run_test(bundle_name, hap_type, experiment_name, scenarios, rounds, target_temperature, trace_tag_mode)
    time.sleep(10)
    analyze(experiment_name)
